package com.baishancloud.log.traffic

import org.apache.flink.api.common.state.{ListState, ListStateDescriptor}
import org.apache.flink.api.common.typeinfo.{TypeHint, TypeInformation}
import org.apache.flink.api.java.utils.ParameterTool
import org.apache.flink.configuration.Configuration
import org.apache.flink.runtime.state.{FunctionInitializationContext, FunctionSnapshotContext}
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction
import org.apache.flink.streaming.api.functions.ProcessFunction
import org.apache.flink.util.Collector

import scala.collection.mutable

/**
 * miniBatch实现，在keyBy之前先对数据进行一部分聚合，以减少下游数据量，进而减少数据倾斜造成的吞吐量降
 *
 * @author ziqiang.wang 
 * @date 2021/12/31 10:40
 */
class MiniBatchProcess(parameterTool: ParameterTool) extends ProcessFunction[LogRecord, LogRecord] with CheckpointedFunction {

  /**
   * 保留checkpoint快照时的中间结果状态
   */
  var listState: ListState[(LogRecordKey, Long)] = _
  /**
   * 保存中间计算结果
   */
  val middleResult: mutable.Map[LogRecordKey, Long] = mutable.Map[LogRecordKey, Long]()
  /**
   * 输出阈值，处理多少条数据之后，输出所有的中间结果，默认值：1w
   */
  var outThreadNum: Long = 10000
  /**
   * 当前处理数据量
   */
  var count: Long = 0

  override def open(parameters: Configuration): Unit = {
    outThreadNum = parameterTool.getLong(outThread, 10000)
  }

  override def processElement(value: LogRecord, ctx: ProcessFunction[LogRecord, LogRecord]#Context, out: Collector[LogRecord]): Unit = {
    val key: LogRecordKey = value.getKey
    if (middleResult.contains(key)) {
      middleResult.put(key, middleResult(key) + value.traffic)
    } else {
      middleResult.put(key, value.traffic)
    }
    count += 1
    if (count % outThreadNum == 0) {
      //  处理数据量达到输出阈值，则输出所有中间结果，然后清空中间结果表
      middleResult.foreach(entry => out.collect(LogRecord(entry._1.timeAt5Min, entry._1.domain, entry._2)))
      middleResult.clear()
    }
  }

  /**
   * 做快照时执行，将中间结果数据写入listState状态
   */
  override def snapshotState(context: FunctionSnapshotContext): Unit = {
    listState.clear()
    middleResult.foreach(entry => listState.add(entry._1, entry._2))
  }

  /**
   * 初始化和恢复快照时执行<br>
   * 初始化listState，将快照数据恢复到中间结果表
   */
  override def initializeState(context: FunctionInitializationContext): Unit = {
    val descriptor: ListStateDescriptor[(LogRecordKey, Long)] = new ListStateDescriptor[(LogRecordKey, Long)]("中间计算结果", TypeInformation.of(new TypeHint[(LogRecordKey, Long)] {}))
    listState = context.getOperatorStateStore.getListState(descriptor)
    if (context.isRestored) {
      middleResult.clear()
      listState.get().forEach(entry => middleResult.put(entry._1, entry._2))
    }
  }
}
